package com.shujia.flink.source

import org.apache.flink.streaming.api.functions.source.SourceFunction
import org.apache.flink.streaming.api.scala._

import java.sql.{Connection, DriverManager, PreparedStatement, ResultSet}

object DemoSourceFunction {
  def main(args: Array[String]): Unit = {
    val env: StreamExecutionEnvironment = StreamExecutionEnvironment.getExecutionEnvironment

    //使用自定义source
    val mysqlDS: DataStream[String] = env.addSource(new MysqlSourceFunction)

    mysqlDS.print()

    env.execute()

  }

  /**
   * 自定义Source 实现SourceFunction接口，实现里面的连个方法
   */
  class MysqlSourceFunction extends SourceFunction[String] {
    /**
     * run只会执行一次，用于生产数据，比如再run方法中使用jdbc读取mysql的数据
     *
     * @param ctx ：上下文对象
     */
    override def run(ctx: SourceFunction.SourceContext[String]): Unit = {
      //1、加载驱动
      Class.forName("com.mysql.jdbc.Driver")
      //2、建立数据库链接
      val con: Connection = DriverManager.getConnection("jdbc:mysql://master:3306/student", "root", "123456")
      //3、编写sql查询数据
      val stat: PreparedStatement = con.prepareStatement("select * from student")
      //4、执行查询
      val resultSet: ResultSet = stat.executeQuery()
      //5、解析数据
      while (resultSet.next()) {
        val id: Int = resultSet.getInt("id")
        val name: String = resultSet.getString("name")
        val age: Int = resultSet.getInt("age")
        val gender: String = resultSet.getString("gender")
        val clazz: String = resultSet.getString("clazz")

        //将数据发送到下游
        ctx.collect(s"$id\t$name\t$age\t$gender\t$clazz")
      }

      //6、关闭链接
      stat.close()
      con.close()
    }
    //在任务取消的时候会被调用，用于回收资源
    override def cancel(): Unit = {}
  }
}
